import torch
import torch.nn as nn
import torch.nn.functional as F

"""
    实现基于负采样的SkipGram模型
"""


class SkipGramModel(nn.Module):
    def __init__(self, vocab_size, embed_size):
        """
        初始化参数
        :param vocab_size: 词表大小
        :param embed_size: 词向量维度
        """
        super(SkipGramModel, self).__init__()
        self.vocab_size = vocab_size
        self.embed_size = embed_size
        # 中心词矩阵
        self.w_embeddings = nn.Embedding(vocab_size, embed_size)
        # 上下文词矩阵
        self.v_embeddings = nn.Embedding(vocab_size, embed_size)
        # 矩阵初始化中心词矩阵、上下文词矩阵
        self._init_emb()

    def _init_emb(self):
        """
        矩阵初始化，word2vec官方提供的初始化方法
        """
        initrange = 0.5 / self.embed_size
        # 中心词矩阵初始化为(-initrange, initrange)的正态分布
        self.w_embeddings.weight.data.uniform_(-initrange, initrange)
        # 周围词矩阵初始化为均值为0的正态分布
        self.v_embeddings.weight.data.uniform_(-0, 0)

    def forward(self, pos_w, pos_v, neg_v):
        """
        前向传播
        :param pos_w:一个batch的中心词
        :param pos_v:一个batch的周围词
        :param neg_v:周围词的负采样，每个中心词应3个周围词的负采样
        :return:返回loss
        """
        # 中心词向量
        # 转为tensor 大小 [ mini_batch_size * emb_dimension ]
        emb_w = self.w_embeddings(torch.LongTensor(pos_w).cuda())
        # 周围词向量
        # 转为tensor 大小 [ mini_batch_size * emb_dimension ]
        emb_v = self.v_embeddings(torch.LongTensor(pos_v).cuda())
        # 转换后大小 [mini_batch_size *  negative_sampling_number * emb_dimension ]
        neg_emb_v = self.v_embeddings(
            torch.LongTensor(neg_v).cuda())
        # 中心词向量与周围词向量对应元素相乘，得到[ mini_batch_size * emb_dimension ]
        score = torch.mul(emb_w, emb_v)
        # 逐列相加得到(batch,1)向量
        score = torch.sum(score, dim=1)
        # 将输入input张量每个元素的截断到区间[min,max]，并返回结果到一个新张量
        score = torch.clamp(score, max=10, min=-10)
        # 先求sigmoid再求log
        score = F.logsigmoid(score)

        # bmm:两个三维矩阵相乘，两个矩阵维度必须相同，得到[mini_batch_size *  negative_sampling_number * 1]
        neg_score = torch.bmm(neg_emb_v, emb_w.unsqueeze(2))
        # 将输入input张量每个元素的截断到区间[min,max]，并返回结果到一个新张量
        neg_score = torch.clamp(neg_score, max=10, min=-10)
        # 负采样的损失函数
        neg_score = F.logsigmoid(-1 * neg_score)
        # L = log sigmoid (Xw.T * θv) + ∑neg(v) [log sigmoid (-Xw.T * θneg(v))]
        # 损失函数越小越好，loss求min
        loss = - torch.sum(score) - torch.sum(neg_score)
        # 返回损失
        return loss

    def save_embedding(self, id2word, file_name):
        """
        保存embedding
        :param id2word: id与单词字典
        :param file_name: 保存的文件名
        """
        # 获取中心词矩阵
        embedding_1 = self.w_embeddings.weight.data.cpu().numpy()
        # 获取周围词矩阵
        embedding_2 = self.v_embeddings.weight.data.cpu().numpy()
        # 两矩阵求平均
        embedding = (embedding_1 + embedding_2) / 2
        fout = open(file_name, 'w')
        fout.write('%d %d\n' % (len(id2word), self.embed_size))
        # 返回键值对
        for wid, w in id2word.items():
            # 向量
            e = embedding[wid]
            e = ' '.join(map(lambda x: str(x), e))
            # 写入格式为词-向量
            fout.write('%s %s\n' % (w, e))


if __name__ == '__main__':
    model = SkipGramModel(100, 10)
    id2word = dict()
    for i in range(100):
        id2word[i] = str(i)
    pos_w = [0, 0, 1, 1, 1]
    pos_v = [1, 2, 0, 2, 3]
    neg_v = [[23, 42, 32], [32, 24, 53], [32, 24, 53], [32, 24, 53], [32, 24, 53]]
    model.forward(pos_w, pos_v, neg_v)
    model.save_embedding(id2word, '../results/test.txt')
